import sys
import xml.parsers.expat
import psycopg2
import pg
import getopt


class OsmData:
    def __init__(self, node_name):
        self.node_name = node_name
        self.tags = dict()
        self.props = dict()
    def clear(self, node_name):
        if (node_name != self.node_name):
            self.node_name = node_name
            self.tags = dict()
            self.props = dict()            
                
                                    
class DBWriter:
    def __init__(self, capacity, prefix=""):
        self._capacity = capacity
        
        global dbhost, dbname, dbuser, dbpass
        self._conn = psycopg2.connect("host='{0}' dbname='{1}' user='{2}' password='{3}'".format(dbhost, dbname, dbuser, dbpass))
        self._cur = self._conn.cursor()
               
        self._pts_filename = prefix + "tmp_nodes"
        self._ways_filename = prefix + "tmp_ways" 
        self._pts_f = open(self._pts_filename, "w")
        self._ways_f = open(self._ways_filename, "w")
        self._count = {"nodes": 0, "ways": 0};        
    
    
    def close_file(self):
        self._pts_f.close()
        self._ways_f.close()
            
    def write_node(self, osm_data):
        self._pts_f.write(str(osm_data.props["id"]) + "|SRID=4326;POINT(" + str(osm_data.props["lon"]) + " " + str(osm_data.props["lat"]) + ")\n")
        self._count["nodes"] += 1
        
                    
    def write_way(self, osm_data):
        nd1 = osm_data.props["nodes"].__iter__()
        nd2 = osm_data.props["nodes"].__iter__()
        nd2.next()
        count = 0
        try:
            while True:
                self._ways_f.write(str(nd1.next()) + "|" + str(nd2.next()) +  "\n")
        except StopIteration:
            pass
        if len(osm_data.props["nodes"]) > 1:
            self._count["ways"] += len(osm_data.props["nodes"]) - 1
                                
        
    def write(self, osm_data):
        if (osm_data.node_name == "node"):
            self.write_node(osm_data)
            if (self._count["nodes"] >= self._capacity):
                self.flush("nodes")
        elif (osm_data.node_name == "way"):
            self.write_way(osm_data)
            if (self._count["ways"] >= self._capacity):
                self.flush("ways")

            
    def flush(self, type):        
        if type == "nodes":
            self._pts_f.flush()
            self._cur.execute("COPY nodes (osm_id, geodata) FROM '{0}/{1}' WITH DELIMITER '|'".format(sys.path[0], self._pts_filename))
            self._pts_f.truncate(0)
            self._pts_f.seek(0)
            self._count["nodes"] = 0
        elif type == "ways":
            self._ways_f.flush()
            self._cur.execute("COPY ways (src, dst) FROM '{0}/{1}' WITH DELIMITER '|'".format(sys.path[0], self._ways_filename))
            self._ways_f.truncate(0)
            self._ways_f.seek(0)
            self._count["ways"] = 0
        self._conn.commit()
        
               
        
class ProgressLoger:
    def __init__(self, f):
        self._file = f
        self._read = 0
        self._untold = 0
        sys.stdout.write("\r  0GB   0MB")
        sys.stdout.flush()

    def read(self, n):
        self._read += n
        self._untold += n
        if self._untold > 10 * 1024 * 1024:
            self.print_read()
        return self._file.read(n)

    def print_read(self):        
        sys.stdout.write("\r{0:>3}GB {1:>3}MB".format(self._read / 1024 / 1024 / 1024, \
                                                      (self._read / 1024 / 1024) % 1024))
        sys.stdout.flush()
        self._untold = 0   
                

def osm_to_db(f):
    writer = DBWriter(2000000)
    
    class Handler:
        def __init__(self, writer):
            self.elem = OsmData("null")
            self.writer = writer        
            self.nodes_done = False
        def start_el(self, name, attrs):
            if name == "node":
                self.elem.clear("node")
                self.elem.props["lat"] = attrs["lat"]
                self.elem.props["lon"] = attrs["lon"]           
                self.elem.props["id"] = attrs["id"]
            elif name == "way":
                if not self.nodes_done:
                    self.writer.flush("nodes")
                    self.nodes_done = True            
                self.elem.clear("way")
                self.elem.props["nodes"] = list()
            elif name == "nd":# and self.elem.node_name == "way":
                self.elem.props["nodes"].append(attrs["ref"])                
        def end_el(self, name):
            if name == "node" or name == "way":
                self.writer.write(self.elem)            

    h = Handler(writer)
    p = xml.parsers.expat.ParserCreate()
    p.StartElementHandler = lambda name, attrs: h.start_el(name, attrs)
    p.EndElementHandler = lambda name: h.end_el(name)
    p.ParseFile(f)
    writer.flush("nodes")
    writer.flush("ways")


dbname = ""
dbhost = ""
dbpass = ""
dbuser = ""

def read_args():
    global dbname, dbhost, dbpass, dbuser
    opts, args = getopt.getopt(sys.argv[1:], "d:u:p:h:", ["dbname=", "user=", "password=", "host="])
    for o, a in opts:
        if o == "-d":
            dbname = a
        elif o == "-u":
            dbuser = a
        elif o == "-p":
            dbpass = a
        elif o == "-h":
            dbhost = a

                
if __name__ == "__main__":
    read_args()
    osm_to_db(ProgressLoger(sys.stdin))
